import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets.mnist import read_image_file, read_label_file

class Data(Dataset):
    # def __init__(self, root, train=True, transform=None, target_transform=None):
    def __init__(self, imagepath, labelpath):
        self.images = read_image_file(imagepath) / 255
        self.labels = read_label_file(labelpath)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        if not isinstance(idx, int):
            raise TypeError
        if idx < 0 or idx >= len(self.labels):
            raise IndexError

        return torch.reshape(self.images[idx], (1, 28, 28)), \
                self.labels[idx]

train_data = Data("./data/MNIST/train-images.idx3-ubyte", "./data/MNIST/train-labels.idx1-ubyte")
test_data = Data("./data/MNIST/t10k-images.idx3-ubyte", "./data/MNIST/t10k-labels.idx1-ubyte")

train_loader = DataLoader(train_data, batch_size=1, shuffle=True, drop_last=True)
test_loader = DataLoader(test_data, batch_size=1, shuffle=True, drop_last=True)